

from types import NoneType
import uvicorn
from fastapi.responses import HTMLResponse
from fastapi import FastAPI, File, UploadFile
from datetime import datetime
import os
from fastapi.responses import FileResponse
from random import randint
from fastapi.middleware.cors import CORSMiddleware
from input_table import create_table_line
from entity import InputFileName,OutputLink,OutputFile
from table_question_line import creat_prompt,to_execl
from wanna import vn



app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/upload-file/")
async def create_upload_file(file: UploadFile = File(...)):
    cp = os.path.join(os.getcwd(),"user_file")
    os.makedirs(cp, exist_ok=True)
    if not file.filename.endswith('.xlsx'):
        return OutputFile(is_err=True,info ="文件必须是.xlsx格式，请重新上传",file_name="")
    file_name  = datetime.now().strftime("%Y%m%d_%H%M%S_") + str(randint(0, 1000)) + ".xlsx"
    try:
        with open(os.path.join(cp,file_name), "wb") as file_object:
            # 读取上传文件的内容
            contents = await file.read()
            # 将内容写入保存文件
            file_object.write(contents)
            # file_object.write(await file.read())
        return OutputFile(is_err=False, info=f"文件 {file.filename} 已上传成功", file_name=file_name)
    except Exception as e:
        return OutputFile(is_err=True,info ="文件上传失败，请重新上传",file_name="")

@app.post("/add_table")
async def add_table(input_data:InputFileName):
    file_name= input_data.file_name
    ans = create_table_line(file_name)
    if ans["state"] == "done":
        return OutputLink(is_err=False, info="数据插入成功", link_line="")
    else:
        return OutputLink(is_err=True, info="数据入库失败", link_line="")

@app.post("/get_ans_file/")
async def ans(input_data:InputFileName):
    file_name = input_data.file_name
    input_str = input_data. user_input
    cp = os.path.join(os.getcwd(), "user_file",file_name)
    prompt,columns = creat_prompt(excel_file=cp,user_input=input_str)
    _, df, _ = vn.ask(question=prompt,auto_train=False,visualize=False,allow_llm_to_see_data=True,print_results=False)
    if isinstance(df,NoneType):
        return OutputLink(is_err=True, info="暂无该相关查询的数据", link_line="")
    else:
        file_name = to_execl(df)
    if isinstance(file_name,NoneType):
        return OutputLink(is_err=True, info="暂无该相关查询的数据", link_line="")
    else:
        link = f'http://192.168.35.221:8053/load_file/{file_name}.xlsx'
        return OutputLink(is_err=False, info="智能填表成功，请下载该表格", link_line=link)


@app.get("/load_file/{file_name}")
async def load_file(file_name:str):
    file_path = os.path.join(os.getcwd(), "out_file",file_name)
    if os.path.exists(file_path):
        return FileResponse(path=file_path,
                        filename=file_name,
                        media_type='application/octet-stream',
                        headers={"Content-Disposition": f'attachment; filename="{file_name}"'})
    else:
        return {"error": "File not found"}, 404

@app.get("/get")
async def main():
    content = """
    <html>
        <body>
            <h1>上传 .xlsx 文件</h1>
            <form action="/upload-file/" enctype="multipart/form-data" method="post">
                <input name="file" type="file" />
                <input type="submit" value="上传文件" />
            </form>
        </body>
    </html>
    """
    return HTMLResponse(content)

uvicorn.run(app=app, host='0.0.0.0', port=8053)